#!/usr/bin/env python3
import csv
import yaml
import numpy as np
import open3d as o3d
from scipy.spatial import KDTree
from sklearn.cluster import DBSCAN
from typing import List, Tuple, Dict

######## 参数 ############
traj_file_path = "../maps/traj.csv"
graph_file_path = "../maps/graph.yaml"
connect_distance_threshold = 0.4  # 距离阈值，用于连接轨迹交叉点
downsample_distance = 1.0         # 降采样距离阈值
used_label = ['x', 'y', 'z']      # CSV文件中使用的列名

def load_trajectory(traj_file_path: str = traj_file_path, 
                   used: List[str] = used_label) -> np.ndarray:
    """
    从CSV文件中加载轨迹数据
    
    参数:
        traj_file_path: 轨迹文件路径
        used: 需要使用的列名列表
        
    返回:
        包含轨迹点的numpy数组 (N, 3)
    """
    traj = []
    with open(traj_file_path, 'r') as f:
        reader = csv.reader(f)
        headers = next(reader)  # 读取标题行
        
        # 获取需要的列索引
        indices = [headers.index(label) for label in used]
        
        for row in reader:
            # 提取指定列的数值并转换为浮点数
            point = [float(row[i]) for i in indices]
            traj.append(point)
    
    return np.array(traj)

def trajectory_filter(traj_list: np.ndarray) -> np.ndarray:
    """
    对轨迹进行滤波，去除异常点
    
    参数:
        traj_list: 输入的轨迹点数组 (N, 3)
        
    返回:
        过滤后的轨迹点数组 (M, 3)
    """
    window_size = 20
    var_threshold = 2.0
    
    # 用于存储每个点的方差
    var_list = np.zeros(len(traj_list))
    
    for i in range(len(traj_list)):
        # 计算窗口范围
        start = max(0, i - window_size // 2)
        end = min(len(traj_list), i + window_size // 2 + 1)
        
        # 提取窗口内的点
        window_points = traj_list[start:end]
        
        # 使用PCA拟合直线
        mean = np.mean(window_points, axis=0)
        cov_matrix = np.cov(window_points - mean, rowvar=False)
        _, eigen_vectors = np.linalg.eigh(cov_matrix)
        direction = eigen_vectors[:, -1]  # 最大特征值对应的方向
        
        # 计算每个点到拟合直线的距离
        diff = window_points - mean
        t = np.dot(diff, direction)
        projected_points = mean + np.outer(t, direction)
        residuals = np.linalg.norm(window_points - projected_points, axis=1)
        
        # 计算窗口内点的方差
        var_list[i] = np.var(residuals)
    
    # 计算全局方差阈值
    mean_var = np.mean(var_list)
    threshold = mean_var * var_threshold
    
    # 过滤异常点
    filtered_indices = np.where(var_list <= threshold)[0]
    return traj_list[filtered_indices]

def create_graph_from_trajectory(traj_list: np.ndarray) -> Dict[str, List]:
    """
    从轨迹生成图结构
    
    参数:
        traj_list: 轨迹点数组 (N, 3)
        
    返回:
        包含节点和边的字典
    """
    # 将每个轨迹点作为节点
    nodes = traj_list.tolist()
    
    # 生成连续边（i -> i+1）
    edges = [(i, i + 1) for i in range(len(nodes) - 1)]
    
    # 使用KDTree查找邻近点
    tree = KDTree(nodes)
    nearby_pairs = tree.query_pairs(connect_distance_threshold)
    
    # 添加邻近点之间的边
    edges.extend(list(nearby_pairs))
    
    # 构建图结构
    graph = {
        "nodes": nodes,
        "edges": edges
    }
    
    # 降采样处理
    graph = graph_down_sample(graph)
    
    return graph

def graph_down_sample(graph: Dict[str, List], 
                     edge_distance_threshold: float = downsample_distance) -> Dict[str, List]:
    """
    图的降采样处理，合并距离过近的节点
    
    参数:
        graph: 输入图结构
        edge_distance_threshold: 距离阈值
        
    返回:
        降采样后的图结构
    """
    nodes = np.array(graph['nodes'])
    edges = graph['edges']
    
    # 使用DBSCAN聚类算法进行降采样
    clustering = DBSCAN(eps=edge_distance_threshold, min_samples=1).fit(nodes)
    labels = clustering.labels_
    
    # 创建新的节点列表（聚类中心）
    unique_labels = np.unique(labels)
    new_nodes = [np.mean(nodes[labels == l], axis=0).tolist() for l in unique_labels]
    
    # 创建节点映射
    label_to_index = {l: i for i, l in enumerate(unique_labels)}
    
    # 生成新的边
    new_edges = []
    seen_edges = set()
    
    for u, v in edges:
        label_u = labels[u]
        label_v = labels[v]
        
        if label_u != label_v:
            new_u = label_to_index[label_u]
            new_v = label_to_index[label_v]
            
            # 避免重复边
            edge_key = tuple(sorted((new_u, new_v)))
            if edge_key not in seen_edges:
                seen_edges.add(edge_key)
                new_edges.append((new_u, new_v))
    
    return {
        "nodes": new_edges,
        "edges": new_edges
    }

def save_graph_file_to_yaml(graph: Dict[str, List], 
                           graph_path: str = graph_file_path):
    """
    将图结构保存为YAML文件
    
    参数:
        graph: 图结构字典
        graph_path: 输出文件路径
    """
    # 将numpy数组转换为列表
    safe_graph = {
        "nodes": [list(map(float, node)) for node in graph['nodes']],
        "edges": [list(map(int, edge)) for edge in graph['edges']]
    }
    
    with open(graph_path, 'w') as f:
        yaml.dump(safe_graph, f, default_flow_style=False)

def graph_visualize(graph: Dict[str, List]):
    """
    使用Open3D可视化图结构
    
    参数:
        graph: 图结构字典
    """
    # 转换为numpy数组
    nodes = np.array(graph['nodes'])
    edges = graph['edges']
    
    # 创建点云
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(nodes)
    
    # 创建线条
    lines = [[u, v] for u, v in edges]
    line_set = o3d.geometry.LineSet(
        points=o3d.utility.Vector3dVector(nodes),
        lines=o3d.utility.Vector2iVector(lines)
    )
    
    # 设置颜色
    line_set.paint_uniform_color([0, 0, 1])  # 蓝色线条
    pcd.paint_uniform_color([1, 0, 0])       # 红色节点
    
    # 可视化
    o3d.visualization.draw_geometries([pcd, line_set], 
                                    window_name="Graph Visualization")

if __name__ == "__main__":
    # 主程序流程
    try:
        # 1. 加载轨迹
        traj_list = load_trajectory()
        print(f"Loaded {len(traj_list)} points from trajectory")
        
        # 2. 轨迹滤波
        traj_list = trajectory_filter(traj_list)
        print(f"Filtered to {len(traj_list)} points")
        
        # 3. 创建图结构
        graph = create_graph_from_trajectory(traj_list)
        print(f"Generated graph with {len(graph['nodes'])} nodes and {len(graph['edges'])} edges")
        
        # 4. 保存图文件
        save_graph_file_to_yaml(graph)
        print(f"Saved graph to {graph_file_path}")
        
        # 5. 可视化
        graph_visualize(graph)
        
    except Exception as e:
        print(f"Error occurred: {str(e)}")